#include <fstream>
#include <iostream>
#include <cxxtest/TestSuite.h>

#include <libgeodecomp.h>
#include <libgeodecomp/geometry/partitions/zcurvepartition.h>
#include <libgeodecomp/io/testinitializer.h>
#include <libgeodecomp/misc/sharedptr.h>
#include <libgeodecomp/misc/testcell.h>
#include <libgeodecomp/parallelization/nesting/mpiupdategroup.h>
#include <libgeodecomp/storage/mockpatchaccepter.h>

using namespace LibGeoDecomp;

namespace LibGeoDecomp {

class MPIUpdateGroupTest : public CxxTest::TestSuite
{
public:
    typedef ZCurvePartition<2> PartitionType;
    typedef VanillaStepper<TestCell<2>, UpdateFunctorHelpers::ConcurrencyNoP> StepperType;
    typedef MPIUpdateGroup<TestCell<2> > UpdateGroupType;
    typedef StepperType::GridType GridType;

    void setUp()
    {
        rank = MPILayer().rank();
        dimensions = Coord<2>(231, 350);

        weights.clear();
        weights << (231 / 2) * (350 * 3 / 4)
                << (231 / 2) * (350 / 2)
                << (231 / 2) * (350 / 4);
        int remainder = dimensions.prod() - sum(weights);
        weights << remainder;

        partition.reset(new PartitionType(Coord<2>(), dimensions, 0, weights));
        ghostZoneWidth = 9;
        init.reset(new TestInitializer<TestCell<2> >(dimensions));
        updateGroup.reset(
            new UpdateGroupType(
                partition,
                CoordBox<2>(Coord<2>(), dimensions),
                ghostZoneWidth,
                init,
                reinterpret_cast<StepperType*>(0)));
        expectedNanoSteps.clear();
        expectedNanoSteps << 5 << 7 << 8 << 33 << 55;
        mockPatchAccepter.reset(new MockPatchAccepter<GridType>());

        for (std::deque<std::size_t>::iterator i = expectedNanoSteps.begin();
             i != expectedNanoSteps.end();
             ++i) {
            mockPatchAccepter->pushRequest(*i);
        }

        updateGroup->addPatchAccepter(mockPatchAccepter, StepperType::INNER_SET);
    }

    void tearDown()
    {
        updateGroup.reset();
    }

    void testUpdate()
    {
        updateGroup->update(100);
        std::deque<std::size_t> actualNanoSteps = mockPatchAccepter->getOfferedNanoSteps();
        TS_ASSERT_EQUALS(actualNanoSteps, expectedNanoSteps);
    }

private:
    std::deque<std::size_t> expectedNanoSteps;
    unsigned rank;
    Coord<2> dimensions;
    std::vector<std::size_t> weights;
    unsigned ghostZoneWidth;
    SharedPtr<PartitionType>::Type partition;
    SharedPtr<Initializer<TestCell<2> > >::Type init;
    SharedPtr<MPIUpdateGroup<TestCell<2> > >::Type updateGroup;
    SharedPtr<MockPatchAccepter<GridType> >::Type mockPatchAccepter;
};

}
